import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import gstools as gs
import json
import yaml
from datetime import datetime
from pykrige.ok import OrdinaryKriging
import matplotlib.colors as mcolors

# 读取数据和配置
with open('data2.json', 'r') as f:
    data_list = json.load(f)
with open('application.yaml', 'r') as f:
    config = yaml.safe_load(f)

# 转换数据为DataFrame
data = pd.DataFrame(data_list, columns=['latitude', 'longitude', 'aqi'])

# 创建更细致的网格以获得更平滑的效果
bbox = config['bounding_box']
gridx = np.linspace(bbox[0], bbox[2], 1)  # 提高分辨率
gridy = np.linspace(bbox[1], bbox[3], 1)

# 优化协方差模型参数
cov_model = gs.Gaussian(dim=2, len_scale=2, anis=0.5, angles=0, var=1.0, nugget=0.1)

# 克里金插值
OK = OrdinaryKriging(
    data['longitude'].values,
    data['latitude'].values,
    data['aqi'].values,
    cov_model
)

# 执行插值
z, ss = OK.execute('grid', gridx, gridy)

# 创建自定义颜色映射
colors = config['colors']
bounds = config['bounds']
# 使用渐变色彩映射
cmap = plt.cm.viridis  # 使用渐变色彩映射
norm = mcolors.Normalize(vmin=min(bounds), vmax=max(bounds))


# 可视化设置
plt.figure(figsize=(14, 10))

# 使用更平滑的等值线绘制
contour = plt.contourf(gridx, gridy, z.T, levels=50, cmap=cmap, alpha=0.9, norm=norm)


# 添加色标
plt.colorbar(contour, label='AQI', ticks=bounds)

# 绘制监测站点
scatter = plt.scatter(
    data['longitude'],
    data['latitude'],
    c=data['aqi'],
    edgecolor='white',
    linewidth=0.5,
    cmap=cmap,
    norm=norm,
    s=30,
    label='监测站'
)

plt.title('中国区域AQI空间插值分布图', pad=20)
plt.xlabel('经度')
plt.ylabel('纬度')
plt.legend(loc='upper right')
plt.grid(True, linestyle='--', alpha=0.3)

# 保存结果
output_path = (config['result_img']['base_path'] +
               datetime.now().strftime(config['result_img']['file_name_template']) +
               config['result_img']['ext_name'])
plt.savefig(output_path, dpi=300, bbox_inches='tight')
plt.show()
